import plotly.express as px
import pandas as pd

df = pd.read_csv('../datasets/all_vaccine_data.csv')
df = df.drop(columns=['DOSES', 'TARGET_NUMBER'])
df = df.dropna()

years_sorted = sorted(df['YEAR'].unique())

fig = px.scatter(df,
                 x='COVERAGE', 
                 y='Life Expectancy', 
                 animation_frame='YEAR', 
                 animation_group='Country Name',
                 size='Child Mortality per 1000 ', 
                 color='continent', 
                 hover_name='Country Name',
                 size_max=55, 
                 range_x=[0, 100], 
                 range_y=[0, 100],
                 height=550,
                 category_orders={'YEAR': years_sorted},
                 custom_data=['Country Name', 'continent', 'Child Mortality per 1000 '])

# Update the hover template for all traces
fig.update_traces(hovertemplate="<b>%{customdata[0]}</b><br>" +
                                "<b>%{customdata[1]}</b><br><br>" +
                                "Vaccination Rate: %{x}%<br>" +
                                "Life Expectancy: %{y} years<br>" +
                                "Child Mortality per 1000: %{customdata[2]}<br>" +
                                "<extra></extra>")

# Update each frame to ensure hover template is consistent
for frame in fig.frames:
    for trace in frame.data:
        trace.hovertemplate = "<b>%{customdata[0]}</b><br>" + \
                              "<b>%{customdata[1]}</b><br><br>" + \
                              "Vaccination Rate: %{x}%<br>" + \
                              "Life Expectancy: %{y} years<br>" + \
                              "Child Mortality per 1000: %{customdata[2]}<br>" + \
                              "<extra></extra>"

fig.update_layout(
    title="Life Expectancy vs Vaccination Rates<br><sup>Child mortality declines as vaccination rate increases</sup>", title_x=0.5,
    plot_bgcolor='#cff8d6',
    paper_bgcolor='#cff8d6',
    margin={'l': 110, 'b': 250, 'r': 130, 't': 100},
    xaxis=dict(title='Vaccination Rate (%)'),   
    yaxis=dict(title='Life Expectancy in years'),
    showlegend=True,
    hovermode='closest',
    
    # Customize the bottom slider
    sliders=[
        {
            'active': 0,
            'yanchor': 'top',
            'xanchor': 'left',

            # Add current year display
            'currentvalue': {
                'font': {'size': 20},
                'prefix': 'Year: ',
                'visible': True,
                'xanchor': 'right'
            },
            'pad': {'b': 10, 't': 50},
            'len': 0.9,
            'x': 0.1,
            'y': 0,
            'steps': [
                {'args': [[year], {'frame': {'duration': 300, 'redraw': True}, 'mode': 'immediate'}],
                 'label': str(year),
                 'method': 'animate'} for year in years_sorted
            ],
            'transition': {'duration': 300},
            'bgcolor': '#fff', 
            'tickcolor': '#000',
        }
    ]
)

# Add caption
fig.add_annotation(x=-0.09, y=-1.12,
                   showarrow=False,
                   xref='paper', yref='paper',
                   xanchor='left', yanchor='bottom',
                   align='left',
                   text='All countries are displayed as bubbles on the graph, where a bigger bubble size corresponds to a higher child mortality.<br>' +
                   'The child mortality has a value type of deaths per 1000 children. The x-axis corresponds to the vaccination in percentages and<br>' +
                   'the y-axis to the vaccination rate of the country. Hovering over a bubble reveals which country the bubble represents.')

fig.show()